package com.ysb.corrector.hmm;

import com.ysb.common.PinyinUtils;
import com.ysb.corrector.common.Item;
import com.ysb.corrector.common.PriopritySet;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @author husc
 * @date 2021/3/19
 */
public class Viterbi {

    public static List<Item> viterbi(AbstractHmmParams hmmParams, String searchKey, int topN, boolean isLogScore, float minProb) {
        Map<String, PriopritySet> preMap = new HashMap<>();
        Map<String, PriopritySet> curMap = new HashMap<>();
        // 转换searchKey为拼音数组
        List<String> observations = PinyinUtils.getPinYin(searchKey);
        // 遍历拼音(观测值)
        List<String> preStates;
        List<String> curStates = Collections.emptyList();
        for (int i = 0; i < observations.size(); i++) {
            curMap = new HashMap<>();
            String observation = observations.get(i);
            // 拼音对应的汉字：observation->state
            preStates = curStates;
            curStates = hmmParams.getStates(observation);

            // 第一个，初始化score
            if (i == 0) {
                for (String state : curStates) {
                    float score = 0.00f;
                    float transitionScore = Math.min(hmmParams.initialProbability(state), minProb);
                    float emissionScore = Math.min(hmmParams.emission(state, observation), minProb);
                    if (isLogScore) {
                        score = (float) Math.log(transitionScore) + (float) Math.log(emissionScore);
                    } else {
                        score = transitionScore + emissionScore;
                    }
                    List<String> path = new LinkedList<>();
                    path.add(state);
                    PriopritySet curMapOrDefault = curMap.getOrDefault(state, new PriopritySet(topN));
                    curMapOrDefault.put(score, path);
                    curMap.put(state, curMapOrDefault);
                }
                preMap = curMap;
                curMap = new HashMap<>();
                continue;
            }

            // 从第2个开始
            for (String state : curStates) {
                System.out.println("");
                for (String preState : preStates) {
                    PriopritySet priopritySet = preMap.get(preState);
                    List<Item> itemList = priopritySet.getItemList();
                    for (Item item : itemList) {
                        float score = 0.00f;
                        float transitionScore = Math.min(hmmParams.transition(preState, state), minProb);
                        float emissionScore = Math.min(hmmParams.emission(state, observation), minProb);
                        if (isLogScore) {
                            score = item.getScore() + (float) Math.log(transitionScore) + (float) Math.log(emissionScore);
                        } else {
                            score = item.getScore() + transitionScore + emissionScore;
                        }
                        List<String> path = new LinkedList<>(item.getPath());
                        path.add(state);
                        PriopritySet curMapOrDefault = curMap.getOrDefault(state, new PriopritySet(topN));
                        curMapOrDefault.put(score, path);
                        curMap.put(state, curMapOrDefault);
                    }
                }
            }
            preMap = curMap;
        }

        // 将curMap中的数据合并，排序，取topN
        List<Item> results = curMap.values().stream().flatMap(p -> p.getItemList().stream())
                .sorted(Comparator.comparing(Item::getScore).reversed()).limit(topN).collect(Collectors.toList());
        return results;
    }

    public static List<Item> viterbi(String searchKey) {
        return viterbi(new DefaultHmmParams(), searchKey, 1, true, 3.14f);
    }

    public static void main(String[] args) {
        System.out.println("viterbi(\"安宫牛黄玩\") = " + viterbi("安宫牛黄玩"));
    }
}
